import gym
import numpy as np

from pyglet_render.viewer import Viewer
from ship_envs.core.position import get_random_position, get_random_target_range_position
from ship_envs.core.ship import Ship
from ship_envs.core.world import World
from ship_envs.core.model_transformer import model2geom
from ship_envs.core import action
from ship_envs.core import observation


class SingleShipEnv(gym.Env):
    def __init__(self):
        super().__init__()
        self.world = World()
        self.world.add_ship(Ship())

        # 执行步数限制为500
        self.step_count = 0
        self.step_count_max = 300
        # 0:输出图片 1：输出速度、位置信息
        self.obs_type = 1

        # spaces
        # 状态空间 此处定义为一个向量和一个坐标点
        # 向量：船->目标点  的向量，x,y范围均为[-1,1]
        # 坐标：船的中心点左边，x,y范围均为[0,1]
        # self.observation_space = gym.spaces.Box(low=np.array([-1, -1, 0, 0]), high=np.array([1, 1, 1, 1]), shape=(4,),
        #                                         dtype=np.float32)
        # 以图像为状态空间
        self.observation_space = observation.get_space(self.obs_type)

        # 此处使用DQN测试环境，设置动作空间为离散，仅控制转向
        # 0: 不转向 1：逆时针1° 2：顺时针1°
        self.action_space = action.get_space()

        # 渲染视图
        self.viewer = None

        self.reward_threshold = 0.1

        self.show_fairway = False



    def get_ship(self):
        # 单智能体场景
        return self.world.get_agent_ship(0)

    def set_show_fairway(self, show):
        self.show_fairway = show

    def step(self, action):
        self.step_count += 1
        self.world.step([action])
        obs = self.observation_callback(self.world, self.get_ship())
        reward = self.reward_callback(self.world, self.get_ship())
        done = self.done_callback(self.world, self.get_ship())
        info = self.info_callback(self.world, self.get_ship())
        return obs, reward, done, info

    def reset(self):
        # 重置船位置及目标点位置
        ship = self.get_ship()
        ship.cur_speed = 0.004
        pos = get_random_position()
        ship.set_pos(pos)
        ship.target.set_pos(get_random_target_range_position(pos, self.distance2targets, self.distance_probs))

        # 重置步数
        self.step_count = 0
        # 刷新obs
        self.world.refresh()
        # 单智能体，obs为整个world的obs  全观测
        return self.observation_callback(self.world, ship)

    def render(self, mode="human"):
        if not self.viewer:
            self.viewer = Viewer(700, 700)
            self.viewer.set_bounds(0, 1, 0, 1)
        self.viewer.geoms = []
        entitys = self.world.get_entitys()
        for entity in entitys:
            self.viewer.add_geom(model2geom(entity.get_show_model()))

        self.viewer.render()

    def seed(self, seed=None):
        np.random.seed(seed)

    def steps_out_range(self):
        return self.step_count > self.step_count_max

    def reward_callback(self, world, ship):
        if self.steps_out_range():
            if ship.dif_with_start_angle_abs() == 0:
                return 1
        if world.has_collision():
            return -10
        if ship.dif_with_start_angle_abs() != 0:
            return -0.001
        return 0

    def observation_callback(self, world, ship):
        if self.obs_type == 0:
            return world.get_obs()
        elif self.obs_type == 1:
            entity_pos = []
            for entity in world.landmarks:
                entity_pos.append(np.array(entity.get_pos(), dtype=float) - np.array(ship.get_pos(), dtype=float))
            # communication of all other agents
            other_pos = []
            for other in world.other_ships:
                if other is ship: continue
                other_pos.append(np.array(other.get_pos(), dtype=float) - np.array(ship.get_pos(), dtype=float))
            # 与初始角度的角度差
            dif_angle = [ship.dif_with_start_angle_abs()]

            # 归一化处理
            dif_angle = np.array(dif_angle) / 90
            # if len(entity_pos) > 0:
            #     entity_pos = np.array(entity_pos) / 500
            # print(other_pos)
            # other_pos = (np.array(other_pos) / 500).tolist()
            return np.concatenate([dif_angle] + entity_pos + other_pos).astype(np.float32)

        return world.get_obs()

    def done_callback(self, world, ship):
        # return world.has_collision() or ship.arrive_target() or ship.out_range() or self.steps_out_range()
        return world.has_collision() or ship.out_range() or ship.dif_with_start_angle_abs() > 90

    def info_callback(self, world, ship):
        return {}
